# -*- coding: utf-8 -*-
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
代码生成上下文管理

提供代码生成过程中的上下文信息管理。
充分利用Pass阶段的结果，不进行重复的类型处理。
"""

from dataclasses import dataclass, field
from typing import Dict, Any, Optional, Set
from ...ir import IRNode, IRModule, IRAssign, IRName
from ..finally_context import FinallyContextStack
from mcpy.symbol_registry import SymbolRegistry


@dataclass
class CodeGenContext:
    """代码生成上下文

    管理代码生成过程中的所有上下文信息。
    核心设计原则：充分利用Pass阶段的结果，不重复处理类型信息。
    """

    # 核心对象
    ir_module: Optional[IRModule] = None
    current_node: Optional[IRNode] = None

    # 生成选项
    in_assignment: bool = False  # 遗留标志，逐步被 expected_cpp_type 取代
    parent_precedence: int = 100
    expected_cpp_type: Optional[str] = None  # 期望的 C++ 类型（用于优化代码生成）

    # 作用域信息
    current_function: Optional[str] = None
    current_class: Optional[str] = None

    # 包含管理
    required_includes: Set[str] = field(default_factory=set)

    # 分层缓存机制
    _global_cache: Dict[str, Any] = field(default_factory=dict)  # 全局共享缓存
    _node_cache: Dict[str, Any] = field(default_factory=dict)  # 节点级别缓存

    # 扩展参数（用于在代码生成过程中传递额外信息）
    extension_params: Dict[str, Any] = field(default_factory=dict)

    # Finally 上下文栈（用于管理嵌套 try/finally 块）
    finally_stack: FinallyContextStack = field(default_factory=FinallyContextStack)

    def create_child_context(self, **kwargs) -> "CodeGenContext":
        """创建子上下文

        Args:
            **kwargs: 要覆盖的属性

        Returns:
            CodeGenContext: 子上下文
        """
        # 复制当前上下文的所有字段
        child_data = {
            "ir_module": self.ir_module,
            "current_node": self.current_node,
            "in_assignment": self.in_assignment,
            "parent_precedence": self.parent_precedence,
            "expected_cpp_type": self.expected_cpp_type,
            "current_function": self.current_function,
            "current_class": self.current_class,
            "required_includes": self.required_includes,  # 共享 includes 集合，不复制
            "_global_cache": self._global_cache,  # 全局缓存共享
            "_node_cache": {},  # 每个子上下文都有独立的节点缓存
            "extension_params": self.extension_params,  # 共享扩展参数
            "finally_stack": self.finally_stack,  # 共享 finally 栈
        }

        # 应用覆盖
        child_data.update(kwargs)

        return CodeGenContext(**child_data)

    def with_node(self, node: IRNode) -> "CodeGenContext":
        """创建带有新节点的上下文

        Args:
            node: 新的当前节点

        Returns:
            CodeGenContext: 新的上下文
        """
        return self.create_child_context(current_node=node)

    def with_precedence(self, precedence: int) -> "CodeGenContext":
        """创建带有新优先级的上下文

        Args:
            precedence: 新的父优先级

        Returns:
            CodeGenContext: 新的上下文
        """
        return self.create_child_context(parent_precedence=precedence)

    def in_assignment_context(self) -> "CodeGenContext":
        """创建赋值上下文

        Returns:
            CodeGenContext: 赋值上下文
        """
        return self.create_child_context(in_assignment=True)

    def require_include(self, include_path: str) -> None:
        """添加所需的包含文件

        Args:
            include_path: 包含文件路径
        """
        if not include_path or not isinstance(include_path, str):
            return

        if include_path.startswith("<") and include_path.endswith(">"):
            include_path = include_path[1:-1]

        self.required_includes.add(include_path)

        if self.ir_module and hasattr(self.ir_module, "require_include"):
            self.ir_module.require_include(include_path)

    # 代码生成阶段直接使用 IR 节点的扩展参数
    def get_cpp_type(self, node: IRNode) -> Optional[str]:
        """获取节点的C++类型

        直接从TypeReplacementPass的结果获取，不进行重复转换。

        Args:
            node: IR节点

        Returns:
            Optional[str]: C++类型，如果无法确定则返回None
        """
        # 直接从TypeReplacementPass的结果获取
        cpp_type = self.get_extension_param_safely(node, "cpp_type")
        if isinstance(cpp_type, str):
            return cpp_type
        return None

    def get_python_type(self, node: IRNode) -> Optional[str]:
        """获取节点的Python类型

        从TypeInferencePass的结果获取。

        Args:
            node: IR节点

        Returns:
            Optional[str]: Python类型，如果无法确定则返回None
        """
        type_info = self.get_extension_param_safely(node, "inferred_type")
        if type_info is None:
            return None

        if hasattr(type_info, "name"):
            return type_info.name
        return None

    def has_type_info(self, node: IRNode) -> bool:
        """检查节点是否有类型信息

        Args:
            node: IR节点

        Returns:
            bool: 是否有类型信息
        """
        return (
            self.get_cpp_type(node) is not None
            or self.get_python_type(node) is not None
        )

    def is_strongly_typed(self, node: IRNode) -> bool:
        """检查节点是否为强类型

        Args:
            node: IR节点

        Returns:
            bool: 是否为强类型
        """
        cpp_type = self.get_cpp_type(node)
        if not cpp_type:
            return False

        # MCPy的弱类型标识
        weak_types = {"mc::variant", "mc::variants", "mc::dict"}

        # 检查前缀匹配
        for weak_type in weak_types:
            if cpp_type == weak_type or cpp_type.startswith(f"{weak_type}<"):
                return False

        return True

    def get_extension_param_safely(
        self, node: IRNode, param_name: str, default: Any = None
    ) -> Any:
        """安全地获取节点的扩展参数

        与直接调用node.get_extension_param不同，此方法会先检查节点是否有该方法，
        避免在节点不支持扩展参数时引发异常。

        Args:
            node: IR节点
            param_name: 参数名
            default: 默认值

        Returns:
            Any: 参数值或默认值
        """
        if not hasattr(node, "get_extension_param"):
            return default

        value = node.get_extension_param(param_name)
        return value if value is not None else default

    # 缓存相关方法
    def cache_set(self, key: str, value: Any) -> None:
        """设置缓存值

        Args:
            key: 缓存键
            value: 缓存值
        """
        self._node_cache[key] = value

    def cache_get(self, key: str, default: Any = None) -> Any:
        """获取缓存值

        Args:
            key: 缓存键
            default: 默认值

        Returns:
            Any: 缓存值，如果没有则返回默认值
        """
        # 先查节点缓存，再查全局缓存
        if key in self._node_cache:
            return self._node_cache[key]
        return self._global_cache.get(key, default)

    def cache_has(self, key: str) -> bool:
        """检查缓存中是否存在指定键

        Args:
            key: 缓存键

        Returns:
            bool: 是否存在
        """
        return key in self._node_cache or key in self._global_cache

    def clear_cache(self) -> None:
        """清空节点缓存"""
        self._node_cache.clear()
        # 注意：不清空全局缓存，因为它可能被其他上下文使用

    def cache_set_global(self, key: str, value: Any) -> None:
        """设置全局缓存值（跨节点共享）

        Args:
            key: 缓存键
            value: 缓存值
        """
        self._global_cache[key] = value

    def cache_set_node(self, key: str, value: Any) -> None:
        """设置节点缓存值（仅当前节点）

        Args:
            key: 缓存键
            value: 缓存值
        """
        self._node_cache[key] = value

    def get_cache_stats(self) -> Dict[str, Any]:
        """获取缓存统计信息

        Returns:
            Dict[str, Any]: 缓存统计信息
        """
        return {
            "global_cache_size": len(self._global_cache),
            "node_cache_size": len(self._node_cache),
            "total_cache_size": len(self._global_cache) + len(self._node_cache),
        }

    # 代码生成阶段的便利方法
    def get_cpp_type_or_auto(self, node: IRNode) -> str:
        """获取节点的 C++ 类型，如果没有则返回 auto

        Args:
            node: IR 节点

        Returns:
            str: C++ 类型字符串，如果没有则返回 "auto"
        """
        cpp_type = self.get_cpp_type(node)
        return cpp_type if cpp_type else "auto"

    def get_loop_var_cpp_type(self, loop_node: IRNode, var_name: str) -> str:
        """获取循环变量的 C++ 类型

        Args:
            loop_node: 循环节点
            var_name: 循环变量名

        Returns:
            str: C++ 类型字符串，如果没有则返回 "auto"
        """
        # 从 TypeReplacementPass 的结果获取
        cpp_type_key = f"loop_var_{var_name}_cpp_type"
        cpp_type = self.get_extension_param_safely(loop_node, cpp_type_key)
        return cpp_type if cpp_type else "auto"

    def get_assignment_target_cpp_type(
        self, assign_node: IRAssign, target_name: str
    ) -> str:
        """获取赋值目标的 C++ 类型

        Args:
            assign_node: 赋值节点
            target_name: 目标变量名

        Returns:
            str: C++ 类型字符串，如果没有则返回 "auto"
        """
        # 尝试从目标表达式获取类型
        target_exprs = assign_node.get_target_exprs()
        for target_expr in target_exprs:
            if isinstance(target_expr, IRName) and target_expr.id == target_name:
                cpp_type = self.get_cpp_type(target_expr)
                if cpp_type:
                    return cpp_type

        # 回退：对于元组解包赋值，不应该从赋值节点获取类型（因为它是 tuple 类型，不是元素类型）
        # 对于元组解包，每个目标变量应该有自己独立的类型
        if assign_node.is_tuple_unpacking:
            return "auto"  # 元组解包时，如果找不到目标变量的类型，返回 auto

        # 回退：尝试从赋值节点本身获取类型（仅用于非元组解包赋值）
        cpp_type = self.get_cpp_type(assign_node)
        return cpp_type if cpp_type else "auto"

    def get_function_return_type(self) -> Optional[str]:
        """获取当前函数的返回类型

        Returns:
            Optional[str]: 函数返回的 C++ 类型，如果无法确定则返回 None
        """
        if not self.current_function or not self.ir_module:
            return None

        # 从 IR 模块中查找函数定义
        from ...ir.declarations.function import IRFunctionNode

        for node in self.ir_module.nodes:
            if isinstance(node, IRFunctionNode) and node.name == self.current_function:
                # 使用函数节点自己的方法获取返回类型
                return node.get_cpp_return_type()

        return None

    # 扩展参数管理方法
    def get_extension_param(self, key: str, default: Any = None) -> Any:
        """获取扩展参数

        Args:
            key: 参数键
            default: 默认值

        Returns:
            Any: 参数值，如果不存在则返回默认值
        """
        return self.extension_params.get(key, default)

    def set_extension_param(
        self, key: str, value: Any, replace_existing: bool = False
    ) -> None:
        """设置扩展参数

        Args:
            key: 参数键
            value: 参数值
            replace_existing: 是否替换已存在的值（默认为 False）
        """
        if not replace_existing and key in self.extension_params:
            return
        self.extension_params[key] = value

    def get_registry(self) -> Optional[SymbolRegistry]:
        """获取符号注册表"""
        return self.ir_module.get_registry() if self.ir_module else None
